from PyQt6.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, 
                             QPushButton, QLabel, QLineEdit, QComboBox,
                             QTextEdit, QFileDialog, QMessageBox, QApplication)
from PyQt6.QtCore import Qt
from PyQt6.QtGui import QPixmap
import sys
import os
import logging
from pathlib import Path

# 获取项目根目录
project_root = Path(__file__).parent.parent.parent
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))

# 导入所需模块
from src.core.api_client import APIClient
from src.utils.config import Config

class MainWindow(QMainWindow):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.api_client = APIClient(config)
        self.setup_ui()
        
    def setup_ui(self):
        self.setWindowTitle("AI Drawing")
        self.setMinimumSize(800, 600)
        
        # 创建中央窗口部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        layout = QVBoxLayout(central_widget)
        
        # 模型选择
        model_layout = QVBoxLayout()
        model_label = QLabel("Model:")
        self.model_combo = QComboBox()
        self.model_combo.setEditable(True)
        self.model_combo.addItem("irag-1.0")
        model_layout.addWidget(model_label)
        model_layout.addWidget(self.model_combo)
        
        # API URL
        url_layout = QVBoxLayout()
        url_label = QLabel("API URL:")
        self.url_combo = QComboBox()
        self.url_combo.setEditable(True)
        self.url_combo.addItem("https://qianfan.baidubce.com/v2/images/generations")
        url_layout.addWidget(url_label)
        url_layout.addWidget(self.url_combo)
        
        # API Key
        key_layout = QVBoxLayout()
        key_label = QLabel("API Key:")
        self.key_input = QLineEdit()
        self.key_input.setEchoMode(QLineEdit.EchoMode.Password)
        key_layout.addWidget(key_label)
        key_layout.addWidget(self.key_input)
        
        # Prompt输入
        prompt_layout = QVBoxLayout()
        prompt_label = QLabel("Prompt:")
        self.prompt_input = QTextEdit()
        self.prompt_input.setMaximumHeight(100)
        prompt_layout.addWidget(prompt_label)
        prompt_layout.addWidget(self.prompt_input)
        
        # 上传图片按钮
        self.upload_btn = QPushButton("Upload Reference Image")
        self.upload_btn.clicked.connect(self.upload_image)
        
        # 生成按钮
        self.generate_btn = QPushButton("Generate")
        self.generate_btn.clicked.connect(self.generate_image)
        
        # 结果显示区域
        result_layout = QVBoxLayout()
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.image_label.setMinimumHeight(300)
        self.text_result = QTextEdit()
        self.text_result.setReadOnly(True)
        self.text_result.setMaximumHeight(100)
        result_layout.addWidget(self.image_label)
        result_layout.addWidget(self.text_result)
        
        # 添加所有组件到主布局
        layout.addLayout(model_layout)
        layout.addLayout(url_layout)
        layout.addLayout(key_layout)
        layout.addLayout(prompt_layout)
        layout.addWidget(self.upload_btn)
        layout.addWidget(self.generate_btn)
        layout.addLayout(result_layout)
        
        # 加载保存的配置
        self.load_config()
        
    def load_config(self):
        """从配置文件加载设置"""
        # TODO: 实现配置加载
        pass
        
    def upload_image(self):
        """处理图片上传"""
        file_name, _ = QFileDialog.getOpenFileName(
            self,
            "Select Image",
            "",
            "Image Files (*.png *.jpg *.jpeg)"
        )
        if file_name:
            # TODO: 处理上传的图片
            logging.info(f"Selected image: {file_name}")
            
    def generate_image(self):
        """处理图片生成"""
        try:
            # 禁用生成按钮，显示正在处理
            self.generate_btn.setEnabled(False)
            self.generate_btn.setText("Generating...")
            self.text_result.setText("Processing request...")
            QApplication.processEvents()  # 确保UI更新
            
            prompt = self.prompt_input.toPlainText().strip()
            if not prompt:
                QMessageBox.warning(self, "Warning", "Please input prompt")
                self.generate_btn.setEnabled(True)
                self.generate_btn.setText("Generate")
                return
                
            model = self.model_combo.currentText().strip()
            url = self.url_combo.currentText().strip()
            api_key = self.key_input.text().strip()
            
            if not api_key:
                QMessageBox.warning(self, "Warning", "Please input API key")
                self.generate_btn.setEnabled(True)
                self.generate_btn.setText("Generate")
                return
                
            if not model:
                QMessageBox.warning(self, "Warning", "Please select or input a model")
                self.generate_btn.setEnabled(True)
                self.generate_btn.setText("Generate")
                return
                
            if not url:
                QMessageBox.warning(self, "Warning", "Please input API URL")
                self.generate_btn.setEnabled(True)
                self.generate_btn.setText("Generate")
                return
                
            # 调用API生成图片
            response = self.api_client.generate_image(
                prompt=prompt,
                model=model,
                api_url=url,
                api_key=api_key
            )
            
            # 显示结果
            if response.get("image"):
                pixmap = QPixmap()
                if pixmap.loadFromData(response["image"]):
                    scaled_pixmap = pixmap.scaled(
                        self.image_label.size(),
                        Qt.AspectRatioMode.KeepAspectRatio,
                        Qt.TransformationMode.SmoothTransformation
                    )
                    self.image_label.setPixmap(scaled_pixmap)
                    self.text_result.setText("Image generated successfully!")
                else:
                    self.text_result.setText("Failed to load generated image.")
            elif response.get("error"):
                self.text_result.setText(f"Error: {response['error']}")
            
            # 显示完整响应信息
            logging.info(str(response.get("info", "")))
            self.text_result.append("\n\nAPI Response:\n" + str(response.get("info", "")))
            
        except Exception as e:
            logging.error(f"Generation error: {str(e)}")
            QMessageBox.critical(self, "Error", f"Generation failed: {str(e)}")
        finally:
            # 恢复生成按钮状态
            self.generate_btn.setEnabled(True)
            self.generate_btn.setText("Generate")